﻿
#include "EulerianSystem.h"
#include "Application.h"
#include "GeneratedMesh.h"
#include "Timer.h"
#include "Transient.h"

#include "Algorithm.h"
#include <vtkm/Matrix.h>
#include <vtkm/TypeTraits.h>
#include <vtkm/cont/ArrayHandleGroupVec.h>
#include <vtkm/worklet/CellMeasure.h>
#include <vtkm/worklet/WorkletCellNeighborhood.h>
#include <vtkm/worklet/WorkletMapTopology.h>

#include <vtkm/exec/ParametricCoordinates.h>
#include <vtkm/worklet/WorkletMapTopology.h>

/// @brief 采用CRTP抽象初始条件的基类，子类实现`ComputeValue`接口
/// @tparam Derived
template<typename Derived>
struct WorkletInitCondition : vtkm::worklet::WorkletVisitCellsWithPoints
{
  using ControlSignature = void(CellSetIn, FieldOut, FieldInPoint);
  using ExecutionSignature = void(CellShape, _2, _3);

  using ValueType = Vec<FloatDefault, 5>;

  template<typename CellShapeTag, typename VType, typename PType>
  VTKM_EXEC void operator()(const CellShapeTag& shape, VType& value, const PType& pts) const
  {
    vtkm::Vec3f lpoint;
    vtkm::exec::ParametricCoordinatesCenter(pts.GetNumberOfComponents(), shape, lpoint);
    typename PType::ComponentType wpoint;
    vtkm::exec::ParametricCoordinatesToWorldCoordinates(pts, lpoint, shape, wpoint);
    auto p = (pts[0] + pts[1] + pts[2] + pts[3]) / 4;
    value = static_cast<const Derived*>(this)->ComputeValue(wpoint);
  }
};


struct CoutteFlowIC : WorkletInitCondition<CoutteFlowIC>
{
  VTKM_EXEC ValueType ComputeValue(const Vec3f& p) const
  {
    FloatDefault gam = 1.4, gamm1 = gam - 1, mach = 0.2;
    FloatDefault u = 1;
    FloatDefault v = 0;
    FloatDefault density = 1;
    FloatDefault pre = 1. / (gam * mach * mach);

    ValueType value;
    value[0] = density;
    value[1] = density * u;
    value[2] = density * v;
    value[3] = density * 0;
    value[4] = pre / gamm1 + 0.5 * density * (u * u + v * v);

    return value;
  }
};


template<UInt8 NUM_VARIABLES>
struct DataPack
{
  VTKM_EXEC
  DataPack()
  {
    for (size_t i = 0; i < NUM_VARIABLES; i++)
    {
      duh[i] = vtkm::TypeTraits<Vec3f>::ZeroInitialization();
    }
  }
  Vec<FloatDefault, NUM_VARIABLES> uh;
  Vec<Vec3f, NUM_VARIABLES> duh;
  //Vec3f duh[NUM_VARIABLES];
  Vec3f inv_term[NUM_VARIABLES]; // 无粘项
  Vec3f vis_term[NUM_VARIABLES]; // 粘性项

  VTKM_EXEC
  void InvisTerm()
  {
    FloatDefault _gamma = 1.4;
    Vec3f vel{ uh[1] / uh[0], uh[2] / uh[0], uh[3] / uh[0] };

    FloatDefault p =
      (_gamma - 1) * (uh[4] - 0.5 * (uh[1] * uh[1] + uh[2] * uh[2] + uh[3] * uh[3]) / uh[0]);
    FloatDefault h = (uh[4] + p) / uh[0];

    inv_term[0] = uh[0] * vel;

    inv_term[1] = uh[1] * vel;
    inv_term[1][0] += p;

    inv_term[2] = uh[2] * vel;
    inv_term[2][1] += p;

    inv_term[3] = uh[3] * vel;
    inv_term[3][2] += p;

    inv_term[4] = uh[0] * h * vel;
  }

  VTKM_EXEC
  void ViscousTerm()
  {
    FloatDefault _reynolds = 10000, _prandtl = 0.72;
    FloatDefault _gamma = 1.4;

    using RealTensor = vtkm::Matrix<FloatDefault, 3, 3>;
    //using RealTensor = Vec<Vec3f, 3>;
    auto r = uh[0];
    auto drho = duh[0];
    Vec3f vel{ uh[1] / uh[0], uh[2] / uh[0], uh[3] / uh[0] };


    RealTensor grad_mom;
    //{ duh[1], duh[2], duh[3] };
    grad_mom[0] = duh[1];
    grad_mom[1] = duh[2];
    grad_mom[2] = duh[3];

    RealTensor grad_vel;
    grad_vel[0] = (grad_mom[0] - vel[0] * drho) / r;
    grad_vel[1] = (grad_mom[1] - vel[1] * drho) / r;
    grad_vel[2] = (grad_mom[2] - vel[2] * drho) / r;

    RealTensor grad_vel_transpose = vtkm::MatrixTranspose(grad_vel);
    FloatDefault vel_div = grad_vel(0, 0) + grad_vel(1, 1) + grad_vel(2, 2);

    RealTensor tau;
    tau[0] = grad_vel[0] + grad_vel_transpose[0];
    tau[1] = grad_vel[1] + grad_vel_transpose[1];
    tau[2] = grad_vel[2] + grad_vel_transpose[2];

    tau(0, 0) -= 2. / 3 * vel_div;
    tau(1, 1) -= 2. / 3 * vel_div;
    tau(2, 2) -= 2. / 3 * vel_div;

    FloatDefault vis = 1;
    for (size_t i = 0; i < 3; i++)
    {
      for (size_t j = 0; j < 3; j++)
      {
        tau(i, j) *= vis / _reynolds;
      }
    }

    Vec3f grad_enthalpy =
      (duh[4] - uh[4] / uh[0] * duh[0]) / r - MatrixMultiply(vtkm::MatrixTranspose(grad_vel), vel);

    for (size_t i = 0; i < 3; i++)
    {
      grad_enthalpy[i] *= (vis / _reynolds) * (_gamma / _prandtl);
    }

    vis_term[0] = vtkm::TypeTraits<Vec3f>::ZeroInitialization();
    vis_term[1] = tau[0];
    vis_term[2] = tau[1];
    vis_term[3] = tau[2];
    vis_term[4] = MatrixMultiply(tau, vel) + grad_enthalpy;
  }
};

struct ComputeResidual : vtkm::worklet::WorkletCellNeighborhood
{
  VTKM_CONT
  ComputeResidual(FloatDefault t)
    : _t(t)
  {
  }

  using ControlSignature =
    void(CellSetIn, FieldInNeighborhood, FieldOut Rhs, FieldIn centroid, FieldIn spacing);
  using ExecutionSignature = void(Boundary, _2, _3, _4, _5);

  template<typename NeighIn, typename RhsType>
  VTKM_EXEC void operator()(const vtkm::exec::BoundaryState& boundary,
                            const NeighIn& value,
                            RhsType& rhs,
                            const Vec3f& centroid,
                            const Vec3f& spacing) const
  {
    //rhs = vtkm::TypeTraits<RhsType>::ZeroInitialization();
    FloatDefault penalty = -8;
    Vec3f normal{ 0, 0, 0 };
    DataPack<RhsType::NUM_COMPONENTS> dl, dr, lift;
    {
      //auto& ul = value.Get(0, 0, 0);
      //auto ur =  value.Get(-1, 0, 0);
      dl.uh = value.Get(0, 0, 0);
      dr.uh = value.Get(-1, 0, 0);
      RhsType lflux, rflux;
      if (boundary.IJK[0] == 0)
      {
        Vec3f p = centroid;
        p[0] = p[0] - 0.5 * spacing[0];
        //ur = ul;
        //dr.uh = value.Get(0, 0, 0);
        dr.uh = valueExactCoutteFlow(_t, p);
      }
      dl.InvisTerm();
      dr.InvisTerm();

      normal = Vec3f{ -1, 0, 0 };
      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; ++i)
      {
        lift.uh[i] = (dl.uh[i] + dr.uh[i]) / 2;
        lift.duh[i] = (dl.uh[i] - dr.uh[i]) / 2 * normal;
      }

      RiemannFlux(dl, dr, normal, lflux);
      lift.ViscousTerm();
      //std::cout << lift.vis_term[0] << std::endl;
      //std::cout << lift.vis_term[1] << std::endl;
      //std::cout << lift.vis_term[2] << std::endl;
      //std::cout << lift.vis_term[3] << std::endl;
      //std::cout << lift.vis_term[4] << std::endl;

      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; i++)
      {
        lflux[i] += penalty * vtkm::Dot(lift.vis_term[i], normal) / spacing[0];
      }


      dl.uh = value.Get(0, 0, 0);
      dr.uh = value.Get(1, 0, 0);
      if (boundary.IJK[0] == boundary.PointDimensions[0] - 1)
      {
        Vec3f p = centroid;
        p[0] = p[0] + 0.5 * spacing[0];
        //ur = ul;
        //dr.uh = value.Get(0, 0, 0);
        dr.uh = valueExactCoutteFlow(_t, p);
      }
      dl.InvisTerm();
      dr.InvisTerm();


      normal = Vec3f{ 1, 0, 0 };
      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; ++i)
      {
        lift.uh[i] = (dl.uh[i] + dr.uh[i]) / 2;
        lift.duh[i] = (dl.uh[i] - dr.uh[i]) / 2 * normal;
      }

      RiemannFlux(dl, dr, normal, rflux);
      lift.ViscousTerm();
      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; i++)
      {
        rflux[i] += penalty * vtkm::Dot(lift.vis_term[i], normal) / spacing[0];
      }

      //RiemannFlux(ul, ur, Vec3f{ 1, 0, 0 }, rflux);
      rhs = (rflux + lflux) / spacing[0];
    }

    {
      //auto& ul = value.Get(0, 0, 0);
      //auto ur = value.Get(0, -1, 0);
      dl.uh = value.Get(0, 0, 0);
      dr.uh = value.Get(0, -1, 0);
      RhsType lflux, rflux;
      if (boundary.IJK[1] == 0)
      {
        Vec3f p = centroid;
        p[1] = p[1] - 0.5 * spacing[1];
        // dr.uh = value.Get(0, 0, 0);
        dr.uh = valueExactCoutteFlow(_t, p);
      }

      //dl.uh = ul;
      //dr.uh = ur;
      dl.InvisTerm();
      dr.InvisTerm();

      normal = Vec3f{ 0, -1, 0 };
      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; ++i)
      {
        lift.uh[i] = (dl.uh[i] + dr.uh[i]) / 2;
        lift.duh[i] = (dl.uh[i] - dr.uh[i]) / 2 * normal;
      }

      RiemannFlux(dl, dr, normal, lflux);
      lift.ViscousTerm();
      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; i++)
      {
        lflux[i] += penalty * vtkm::Dot(lift.vis_term[i], normal) / spacing[0];
      }

      //ur = value.Get(0, 1, 0);
      dl.uh = value.Get(0, 0, 0);
      dr.uh = value.Get(0, 1, 0);
      if (boundary.IJK[1] == boundary.PointDimensions[1] - 1)
      {
        Vec3f p = centroid;
        p[1] = p[1] + 0.5 * spacing[1];
        //dr.uh = value.Get(0, 0, 0);
        dr.uh = valueExactCoutteFlow(_t, p);
      }
      dl.InvisTerm();
      dr.InvisTerm();

      normal = Vec3f{ 0, 1, 0 };
      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; ++i)
      {
        lift.uh[i] = (dl.uh[i] + dr.uh[i]) / 2;
        lift.duh[i] = (dl.uh[i] - dr.uh[i]) / 2 * normal;
      }

      RiemannFlux(dl, dr, normal, rflux);
      lift.ViscousTerm();
      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; i++)
      {
        rflux[i] += penalty * vtkm::Dot(lift.vis_term[i], normal) / spacing[0];
      }

      rhs += (rflux + lflux) / spacing[1];
    }
  }

  VTKM_EXEC
  vtkm::Vec<FloatDefault, 5> valueExactIsoVortex(FloatDefault t, const Vec3f& p) const
  {
    Vec3f _vel{ 1, 1, 0 };
    FloatDefault x = p[0] - _vel[0] * t;
    FloatDefault y = p[1] - _vel[1] * t;
    FloatDefault z = p[2] - _vel[2] * t;
    FloatDefault gam = 1.4, gamm1 = gam - 1, epi = 5.;
    FloatDefault xb, yb, r2;

    FloatDefault PI = vtkm::Pi();
    xb = x + 5;
    yb = y + 5;
    r2 = xb * xb + yb * yb;

    FloatDefault u = _vel[0] + epi / (2. * PI) * exp(0.5 * (1. - r2)) * (-yb);
    FloatDefault v = _vel[1] + epi / (2. * PI) * exp(0.5 * (1. - r2)) * xb;
    FloatDefault T = 1. - gamm1 * epi * epi / (8 * gam * PI * PI) * exp(1. - r2);
    FloatDefault density = pow(T, 1 / gamm1);
    FloatDefault pre = pow(density, gam);

    vtkm::Vec<FloatDefault, 5> value;
    value[0] = density;
    value[1] = density * u;
    value[2] = density * v;
    value[3] = density * 0;
    value[4] = pre / gamm1 + 0.5 * density * (u * u + v * v);

    return value;
  }

  VTKM_EXEC
  vtkm::Vec<FloatDefault, 5> valueExactCoutteFlow(FloatDefault t, const Vec3f& p) const
  {
    FloatDefault _gamma = 1.4, _prandtl = 0.72, _mach = 0.2;

    FloatDefault x = p[0];
    FloatDefault y = p[1];
    FloatDefault z = p[2];

    FloatDefault T2 = 0.85, T1 = 0.8;
    FloatDefault tem =
      T1 + (T2 - T1) * y / 2 + 0.5 * _prandtl * (_gamma - 1) * _mach * _mach * y / 2 * (1 - y / 2);
    FloatDefault pre = 1. / (_gamma * _mach * _mach);
    FloatDefault density = pre * _gamma * _mach * _mach / tem;
    FloatDefault u = y / 2.0, v = 0, w = 0;
    Vec3f mom{ density * u, density * v, density * w };
    FloatDefault rhoe = pre / (_gamma - 1) + 0.5 * density * (u * u + v * v + w * w);

    vtkm::Vec<FloatDefault, 5> value;
    value[0] = density;
    value[1] = density * u;
    value[2] = density * v;
    value[3] = density * w;
    value[4] = rhoe;

    return value;
  }
  template<typename DataPackType, typename VType>
  VTKM_EXEC void RiemannFlux(const DataPackType& dl,
                             const DataPackType& dr,
                             const Vec3f& normal,
                             VType& riemann_flux) const
  {
    for (IdComponent eq = 0; eq < VType::NUM_COMPONENTS; eq++)
    {
      riemann_flux[eq] =
        vtkm::Dot(dl.inv_term[eq] + dr.inv_term[eq], normal) / 2 + 0.6 * (dl.uh[eq] - dr.uh[eq]);
    }
  }
  FloatDefault _t;
};

struct ExplictFirtOrder : public vtkm::worklet::WorkletMapField
{
  VTKM_CONT
  ExplictFirtOrder(FloatDefault dt)
    : _dt(dt)
  {
  }

  using ControlSignature = void(FieldInOut var, FieldIn rhs);
  using ExecutionSignature = void(_1, _2);

  template<typename T>
  VTKM_EXEC void operator()(T& u, const T& rhs) const
  {
    u += -_dt * rhs;
  }

  FloatDefault _dt;
};

struct WriteVariable : public vtkm::worklet::WorkletMapField
{

  using ControlSignature = void(FieldIn solution,
                                FieldOut density,
                                FieldOut velocity,
                                FieldOut rhoe);
  using ExecutionSignature = void(_1, _2, _3, _4);

  template<typename SolType>
  VTKM_EXEC void operator()(const SolType& sol, Real& density, Vec3f& momnetum, Real& rhoe) const
  {
    density = sol[0];
    momnetum = Vec3f{ sol[1], sol[2], sol[3] };
    rhoe = sol[4];
  }
};


RegisterObject(EulerianSystem);
EulerianSystem::EulerianSystem(const Configuration& cfg)
  : System(cfg)
{
}

void EulerianSystem::Init()
{
  _data_set.SetCellSet(_mesh.GetCellSet());
  _data_set.AddCoordinateSystem(_mesh.GetCoordinate());

  // 通过初始条件的赋值变量
  this->AddCellVariable<Real>("density");
  this->AddCellVariable<Vec3f>("momentum");
  this->AddCellVariable<Real>("rhoe");

  // 对 density, velocity, rhoe赋值
  this->InitialCondition();

  // 写回DataSet
  auto density = GetVariableArray<Real>("density");
  auto momentum = GetVariableArray<Vec3f>("momentum");
  auto rhoe = GetVariableArray<Real>("rhoe");
  Invoker{}(WriteVariable{}, _solution, density, momentum, rhoe);

  Info("EulerianSystem::Init");
}

void EulerianSystem::InitialCondition()
{
  vtkm::cont::Timer timer;
  timer.Start();
  vtkm::cont::Invoker invoker;
  invoker(CoutteFlowIC{}, _mesh.GetCellSet(), _solution, _mesh.GetCoordinate());

  Success("初始条件设置完成", timer.GetElapsedTime(), "s");
}

void EulerianSystem::Evolve()
{
  Timer timer;
  timer.Start();

  PreSolve();
  Solve();
  PostSolve();

  Info("EulerianSystem::Evolve运行完成，", "时间: ", timer.GetElapsedTime(), " s");
}

void EulerianSystem::PreSolve() {}

void EulerianSystem::Solve()
{
  static Real _t = 0;
  auto& transient = static_cast<Transient&>(*(_app.GetExecutioner()));
  Real _dt = transient.Dt();

  auto& mesh = static_cast<GeneratedMesh&>(_mesh);
  vtkm::cont::Invoker{}(
    ComputeResidual{ _t }, _mesh.GetCellSet(), _solution, _rhs, mesh._centroid, mesh._spacing);
  //fmt::print(" ElapsedTime: {0:.3f}\n", timer.GetElapsedTime());

  //timer.Start();
  vtkm::cont::Invoker{}(ExplictFirtOrder{ _dt }, _solution, _rhs);
  auto result = vtkm::cont::Algorithm::Reduce(_rhs, vtkm::TypeTraits<VType>::ZeroInitialization());
  std::cout << "Res: " << result[0] / _rhs.GetNumberOfValues() << std::endl;
  _t += _dt;
}

void EulerianSystem::PostSolve()
{
  auto density = GetVariableArray<Real>("density");
  auto momentum = GetVariableArray<Vec3f>("momentum");
  auto rhoe = GetVariableArray<Real>("rhoe");

  Invoker{}(WriteVariable{}, _solution, density, momentum, rhoe);
}
